#include <rockchip/mpp_frame.h>
#include <rockchip/rk_mpi.h>
#include <stdio.h>

#include <opencv4/opencv2/opencv.hpp>

#define MPP_ALIGN(x, a) (((x) + (a)-1) & ~((a)-1))

#define SZ_1K (1024)
#define SZ_2K (SZ_1K * 2)
#define SZ_4K (SZ_1K * 4)

int main(int argc, char* argv[]) {
  // 配置
  MppCtx ctx = nullptr;
  MppApi* mpi = nullptr;
  int cfg_frame_width = 1280;
  int cfg_frame_height = 768;
  int cfg_hor_stride = MPP_ALIGN(cfg_frame_width, 16);
  int cfg_ver_stride = MPP_ALIGN(cfg_frame_height, 16);
  MppFrameFormat cfg_format = MPP_FMT_YUV420P;
  MppCodingType cfg_type = MPP_VIDEO_CodingAVC;  // h264
  int cfg_bps = 600000;      // CBR 模式下的目标码率。
  int cfg_bps_min = 400000;  // VBR 模式下的最高码率。
  int cfg_bps_max = 800000;  // VBR 模式下的最低码率。
  int cfg_rc_mode =
      MppEncRcMode::MPP_ENC_RC_MODE_VBR;  // 码率控制模式:CBR 为 Constant Bit
                                          // Rate，固定码率模式,VBR 为 Variable
                                          // Bit Rate，可变码率模式.
  int cfg_frame_num = 0;

  int cfg_gop_mode = 0;
  int cfg_gop_len = 0;
  int cfg_vi_len = 0;
  int cfg_fps_in_flex = 0;
  int cfg_fps_in_den = 1;
  int cfg_fps_in_num = 8;  // 输入的fps
  int cfg_fps_out_flex = 0;
  int cfg_fps_out_den = 1;
  int cfg_fps_out_num = 8;
  int cfg_frame_size = 0;
  int cfg_header_size = 0;
  int cfg_quiet = 0;

  MppBufferGroup buf_grp;
  MppBuffer frm_buf;
  MppBuffer pkt_buf;
  MppPollType timeout = MPP_POLL_BLOCK;
  MppEncCfg cfg;

  // 设置frame的缓冲大小
  switch (cfg_format) {
    case MPP_FMT_YUV420SP:
    case MPP_FMT_YUV420P: {
      cfg_frame_size =
          MPP_ALIGN(cfg_hor_stride, 64) * MPP_ALIGN(cfg_ver_stride, 64) * 3 / 2;
      printf("frame size %d\n", cfg_frame_size);
    } break;

    default: {
      printf("format not support\n");
      return -1;
    } break;
  }

  //
  if (MPP_FRAME_FMT_IS_FBC(cfg_format)) {
    cfg_header_size = MPP_ALIGN(
        MPP_ALIGN(cfg_frame_width, 16) * MPP_ALIGN(cfg_frame_height, 16) / 16,
        SZ_4K);
  } else {
    cfg_header_size = 0;
  }
  printf("cfg_header_size %d\n", cfg_header_size);

  // 申请缓冲区
  int ret = mpp_buffer_group_get_internal(&buf_grp, MPP_BUFFER_TYPE_DRM);
  if (MPP_OK != ret) {
    printf("failed to get mpp buffer group ret %d\n", ret);
  }

  ret = mpp_buffer_get(buf_grp, &frm_buf, cfg_frame_size + cfg_header_size);
  if (MPP_OK != ret) {
    printf("failed to get frame buffer ret %d\n", ret);
  }

  ret = mpp_buffer_get(buf_grp, &pkt_buf, cfg_frame_size);
  if (MPP_OK != ret) {
    printf("failed to get packet buffer ret %d\n", ret);
  }

  // 创建接口
  ret = mpp_create(&ctx, &mpi);
  if (MPP_OK != ret) {
    printf("mpp create failed \n");
    return -1;
  }

  // 设置接口参数
  ret = mpi->control(ctx, MPP_SET_OUTPUT_TIMEOUT, &timeout);

  // 初始化接口
  ret = mpp_init(ctx, MppCtxType::MPP_CTX_ENC,
                 MppCodingType ::MPP_VIDEO_CodingAVC);
  if (MPP_OK != ret) {
    printf("mpp init failed \n");
    return -1;
  }

  // 初始化编码配置
  ret = mpp_enc_cfg_init(&cfg);
  if (ret) {
    printf("mpp init cfg failed \n");
    return -1;
  }

  mpp_enc_cfg_set_s32(cfg, "prep:width", cfg_frame_width);
  mpp_enc_cfg_set_s32(cfg, "prep:height", cfg_frame_height);
  mpp_enc_cfg_set_s32(cfg, "prep:hor_stride", cfg_hor_stride);
  mpp_enc_cfg_set_s32(cfg, "prep:ver_stride", cfg_ver_stride);
  mpp_enc_cfg_set_s32(cfg, "prep:format", cfg_format);

  mpp_enc_cfg_set_s32(cfg, "rc:mode", cfg_rc_mode);

  //
  mpp_enc_cfg_set_s32(cfg, "rc:fps_in_flex", cfg_fps_in_flex);
  mpp_enc_cfg_set_s32(cfg, "rc:fps_in_num", cfg_fps_in_num);
  mpp_enc_cfg_set_s32(cfg, "rc:fps_in_denorm", cfg_fps_in_den);
  mpp_enc_cfg_set_s32(cfg, "rc:fps_out_flex", cfg_fps_out_flex);
  mpp_enc_cfg_set_s32(cfg, "rc:fps_out_num", cfg_fps_out_num);
  mpp_enc_cfg_set_s32(cfg, "rc:fps_out_denorm", cfg_fps_out_den);
  mpp_enc_cfg_set_s32(cfg, "rc:gop",
                      cfg_gop_len ? cfg_gop_len : cfg_fps_out_den * 2);

  // drop frame or not when bitrate overflow
  mpp_enc_cfg_set_s32(cfg, "rc:drop_mode", MPP_ENC_RC_DROP_FRM_DISABLED);
  mpp_enc_cfg_set_s32(cfg, "rc:drop_thd", 20);  // 20% of max bps
  mpp_enc_cfg_set_s32(cfg, "rc:drop_gap", 1);   // Do not continuous drop frame

  mpp_enc_cfg_set_s32(cfg, "rc:bps_target", cfg_bps);

  switch (cfg_rc_mode) {
    case MPP_ENC_RC_MODE_VBR:
    case MPP_ENC_RC_MODE_AVBR: {
      mpp_enc_cfg_set_s32(cfg, "rc::bps_max",
                          cfg_bps_max ? cfg_bps_max : cfg_bps * 17 / 16);
      mpp_enc_cfg_set_s32(cfg, "rc::bps_min",
                          cfg_bps_min ? cfg_bps_min : cfg_bps * 1 / 16);
    } break;

    default: {
      printf("error cfg rc mode\n");
      return -1;
    } break;
  }

  //
  switch (cfg_type) {
    case MPP_VIDEO_CodingAVC:
    case MPP_VIDEO_CodingHEVC: {
      switch (cfg_rc_mode) {
        case MPP_ENC_RC_MODE_VBR:
        case MPP_ENC_RC_MODE_AVBR: {
          mpp_enc_cfg_set_s32(cfg, "rc:qp_init", 26);
          mpp_enc_cfg_set_s32(cfg, "rc:qp_max", 51);
          mpp_enc_cfg_set_s32(cfg, "rc:qp_min", 10);
          mpp_enc_cfg_set_s32(cfg, "rc:qp_max_i", 51);
          mpp_enc_cfg_set_s32(cfg, "rc:qp_min_i", 10);
          mpp_enc_cfg_set_s32(cfg, "rc:qp_ip", 2);

        } break;

        default: {
          printf("error cfg type \n");
          return -1;
        }

        break;
      }
    }
    /* code */
    break;

    default: {
      printf("error cfg type \n");
      return -1;
    } break;
  }

  // setup codec
  mpp_enc_cfg_set_s32(cfg, "codec:type", cfg_type);
  switch (cfg_type) {
    case MPP_VIDEO_CodingAVC: {
      /*
       *H.264 profile_idc parameter
       *66  - Baseline profile
       *77  - Main profile
       *100 - High profile
       */
      mpp_enc_cfg_set_s32(cfg, "h264:profile", 66);

      /*
       *H.264 level_idc parameter
       *10 / 11 / 12 / 13    - qcif@15fps / cif@7.5fps / cif@15fps / cif@30fps
       *20 / 21 / 22         - cif@30fps / half-D1@@25fps / D1@12.5fps
       *30 / 31 / 32         - D1@25fps / 720p@30fps / 720p@60fps
       *40 / 41 / 42         - 1080p@30fps / 1080p@30fps / 1080p@60fps
       *50 / 51 / 52         - 4K@30fps
       */
      mpp_enc_cfg_set_s32(cfg, "h264:level", 40);
      mpp_enc_cfg_set_s32(cfg, "h264:cabac_en", 1);
      mpp_enc_cfg_set_s32(cfg, "h264:cabac_idc", 0);
      mpp_enc_cfg_set_s32(cfg, "h264:trans8x8", 0);

    } break;

    default: {
      printf("error type\n");
      return -1;
    } break;
  }

  RK_U32 cfg_split_mode = 0;
  RK_U32 cfg_split_arg = 0;
  ret = mpi->control(ctx, MPP_ENC_SET_CFG, cfg);
  if (ret) {
    printf("mpi control enc set cfg failed ret %d\n", ret);
    return -1;
  }

  MppEncSeiMode cfg_sei_mode = MPP_ENC_SEI_MODE_ONE_FRAME;
	ret = mpi->control(ctx, MPP_ENC_SET_SEI_CFG, &cfg_sei_mode);
	if (ret)
	{
		printf("mpi control enc set sei cfg failed ret %d\n", ret);
        return -1;
	}



	if (cfg_type == MPP_VIDEO_CodingAVC || cfg_type == MPP_VIDEO_CodingHEVC)
	{
		MppEncHeaderMode cfg_header_mode = MPP_ENC_HEADER_MODE_EACH_IDR;
		ret = mpi->control(ctx, MPP_ENC_SET_HEADER_MODE, &cfg_header_mode);
		if (ret)
		{
			printf("mpi control enc set header mode failed ret %d\n", ret);
		}
	}

    // 没有设置
    RK_U32 gop_mode = cfg_gop_mode;

    // 省略了一段... setup test mode by env


    FILE *fp_output {nullptr};
    fp_output = fopen("./sample.h264", "wb");

    if(cfg_type == MPP_VIDEO_CodingAVC){
        MppPacket packet = NULL;
        mpp_packet_init_with_buffer(&packet, pkt_buf);
        mpp_packet_set_length(packet, 0);

        ret = mpi->control(ctx, MPP_ENC_GET_HDR_SYNC, packet);
		if (ret)
		{
			printf("mpi control enc get extra info failed\n");
            return -1;
		}
		else
		{ /*get and write sps/pps for H.264 */

			void *ptr = mpp_packet_get_pos(packet);
			size_t len = mpp_packet_get_length(packet);

			//if (p->fp_output)
			fwrite(ptr, 1, len, fp_output);
		}

        mpp_packet_deinit(&packet);
    }

    cv::Mat bgr_frame = cv::imread("bird.jpg");
    cv:resize(bgr_frame, bgr_frame, cv::Size(cfg_frame_width, cfg_frame_height), 0, 0, cv::INTER_LINEAR);
    printf("MPP_ALIGN(cfg_hor_stride, 64): %d\n", MPP_ALIGN(cfg_hor_stride, 64));
    printf("MPP_ALIGN(cfg_ver_stride, 64): %d\n", MPP_ALIGN(cfg_ver_stride, 64));
    printf("start ..\n");

    for( int frame_id = 0; frame_id < 100; frame_id++){
        MppMeta meta = nullptr;
        MppFrame frame = nullptr;
        MppPacket packet = nullptr;

        RK_S32 cam_frm_idx = -1;
        MppBuffer cam_buf = nullptr;
        RK_U32 eoi = 1;
        cv::Mat yuv_frame;
        cv::cvtColor(bgr_frame, yuv_frame, cv::COLOR_BGR2YUV_I420);

        // printf("yuv_frame width:%d. height:%d \n", bgr_frame.cols,
        //        bgr_frame.rows);

        // printf("yuv_frame width:%d. height:%d \n", yuv_frame.cols,
        //        yuv_frame.rows);

        ret = mpp_frame_init(&frame);
        if(ret){
            printf("frame init failed \n");
            return -1;
        }

        mpp_frame_set_width(frame, cfg_frame_width);
        mpp_frame_set_height(frame, cfg_frame_height);
        mpp_frame_set_hor_stride(frame, cfg_hor_stride);
        mpp_frame_set_ver_stride(frame, cfg_ver_stride);
        mpp_frame_set_fmt(frame, cfg_format);
        mpp_frame_set_eos(frame, 0);

        ret = mpp_buffer_write(frm_buf, 0, yuv_frame.data, cfg_frame_size);
        if(ret){
            printf("could not write data on frame\n");
            return -1;
        }

        mpp_frame_set_buffer(frame, frm_buf);
        meta = mpp_frame_get_meta(frame);
        mpp_packet_init_with_buffer(&packet, pkt_buf);

        mpp_packet_set_length(packet, 0);
        mpp_meta_set_packet(meta, KEY_OUTPUT_PACKET, packet);

        ret = mpi->encode_put_frame(ctx, frame);
        if(ret){
             printf("encode_put_frame failed \n");
            return -1;
        }
        mpp_frame_deinit(&frame);
        
        do{
            ret = mpi->encode_get_packet(ctx, &packet);
            if(ret){
                 printf("encode_get_packet failed\n");
                 return -1;
            }
            
            if(packet){
                void *ptr = mpp_packet_get_pos(packet);
                size_t len = mpp_packet_get_length(packet);
                fwrite(ptr, 1, len, fp_output);
                if(mpp_packet_is_partition(packet)){
                    eoi = mpp_packet_is_eoi(packet);
                }
            }
            else{
                printf("there is no packet\n");
            }

            mpp_packet_deinit(&packet);

        } while(!eoi);
    }

    if(ctx){
        mpp_destroy(ctx);
        ctx = nullptr;
    }

    if(cfg) {
        mpp_enc_cfg_deinit(cfg);
        cfg = nullptr;
    }

    fclose(fp_output);

    printf("running successfully\n");
  return 0;
}